/**
 * Copyright (c) Huawei Technologies Co., Ltd. 2023. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#ifndef AIR_CXX_RUNTIME_V2_OP_IMPL_DYNAMIC_GATHERV2_H_
#define AIR_CXX_RUNTIME_V2_OP_IMPL_DYNAMIC_GATHERV2_H_

#include <cstdint>
#include <vector>
#include "op_tiling.h"

namespace TbeReduce {
    constexpr int32_t HALF_UB = 2;
    constexpr int32_t DATA_VALUE = 1024;
    constexpr int32_t NUM_32 = 32;
    constexpr int32_t ACTUAL_NUM = 56.5;
    constexpr int32_t GATE_VALUE = 0.012;
    const int64_t GATHER_BLOCK_SIZE = 32;
    const int64_t PARAMS_CACHED_UB = 100 * 1024;
    const int64_t INDICES_MIN_NUM_FOR_CACHE = 10 * 1024;
    const int64_t RESERVED_UB_SIZE = 8 * 1024;
    const int64_t CACHE_MODE_UB_SLICE = 6;
    const int64_t NOT_ALIGN_MAX_PARAM_ROW = 64;
    const int64_t NOT_ALIGN_MIN_PARAM_ROW = 8;
    const int64_t NOT_ALIGN_SUPPORT_DATA_SIZE = 4;
    const int64_t RESERVED_UB_SIZE_2K = 2 * 1024;
    const int64_t TRANS_POSE_LINE_SIZE = 16;
    const int64_t SUPPORT_PARAM_SIZE = 2;
    const int64_t THREE_PART_UB_SIZE = 3;
    const int64_t FOUR_PART_UB_SIZE = 4;
    // A. block tiling: indices tiling
    // 1. one params row size is smaller than 32B
    // params is not cache
    const int64_t TILING_MODE_1 = 1;
    // params is cache in UB
    const int64_t TILING_MODE_4 = 4;
    // params is cache in L1
    const int64_t TILING_MODE_13 = 13;

    // 2. one params row size is greater than or equal to 32B
    // params_row is not 32B aligned
    const int64_t TILING_MODE_2 = 2;
    // the data of one params row can not store in half UB, need tiling
    const int64_t TILING_MODE_5 = 5;

    // 3. params_row is 32B aligned
    // params is not cache in UB or L1
    const int64_t TILING_MODE_3 = 3;
    // params is cache in UB
    const int64_t TILING_MODE_6 = 6;
    // params is cache in L1
    const int64_t TILING_MODE_7 = 7;

    // B. block tiling: params_pre tiling
    // 1. one params row size is smaller than 32B
    // params is not cache
    const int64_t TILING_MODE_8 = 8;
    // params is cache in UB
    const int64_t TILING_MODE_9 = 9;

    // 2. params_row is 32B aligned
    // params is not cache in UB or L1
    const int64_t TILING_MODE_10 = 10;
    // params is cache in UB
    const int64_t TILING_MODE_11 = 11;
    // params is cache in L1
    const int64_t TILING_MODE_12 = 12;

    // tiling_mode with impl_mode
    const int64_t TILING_MODE_14 = 14;

    // basing TILING_MODE_14 gatherv2 function module,
    // sampling indices and sorting topN for cache indices
    const int64_t TILING_MODE_15 = 15;
    const int64_t TILING_MODE_16 = 16;
    const int64_t TILING_MODE_17 = 17;
    const int64_t TILING_MODE_18 = 18;

    // tiling_mode with batch_dims
    // 1.one params row size is smaller than 32B
    // 1.1 params is cached in UB
    const int64_t TILING_MODE_20 = 20;
    const int64_t TILING_MODE_21 = 21;
    const int64_t TILING_MODE_22 = 22;
    // 1.2 params is not cached in UB
    const int64_t TILING_MODE_23 = 23;
    const int64_t TILING_MODE_24 = 24;
    const int64_t TILING_MODE_25 = 25;

    // 2.one params row size is large than 32B and not align
    const int64_t TILING_MODE_26 = 26;
    const int64_t TILING_MODE_27 = 27;
    const int64_t TILING_MODE_28 = 28;

    // 3.one params row size is align
    // 3.1 params is cached in UB
    const int64_t TILING_MODE_29 = 29;
    const int64_t TILING_MODE_30 = 30;
    const int64_t TILING_MODE_31 = 31;
    // 3.2 params is not cached in UB
    const int64_t TILING_MODE_32 = 32;
    const int64_t TILING_MODE_33 = 33;
    const int64_t TILING_MODE_34 = 34;

    // 4. large params row size
    const int64_t TILING_MODE_35 = 35;
    const int64_t TILING_MODE_36 = 36;
    const int64_t TILING_MODE_37 = 37;

    // 5. small indices row size
    const int64_t TILING_MODE_38 = 38;
    const int64_t TILING_MODE_39 = 39;

    // 6. small params and indices row size
    const int64_t TILING_MODE_40 = 40;
    const int64_t TILING_MODE_41 = 41;

    // define batch_dims of gather_v2 attr idx and name
    // const struct ops::AttrBase GATHER_V2_BATCH_DIMS_INFO(0, "batch_dims");
    // define batch_dims of gather attr idx and name
    // const struct ops::AttrBase GATHER_BATCH_DIMS_INFO(1, "batch_dims");
    static const int64_t BATCH_DIMS_DEFAULT_VALUE = 0;

    // define impl_mode of gather_v2 attr
    static const int64_t IMPL_MODE_DEFAULT_VALUE = 0;
    static const int64_t IMPL_MODE_HIGH_PERFORMANCE_VALUE = 1;

    // define cache mode gather_v2 param
    static const int64_t TILING_MODE_15_MAX_PARAM_NUM_SIZE = 1024;


    struct GatherCompileInfos {
        // std::shared_ptr<AutoTilingHandler> outer_compile_info;
        int64_t ub_size{1};
        int64_t l1_size{0};
        int64_t core_num{1};
        int64_t params_d_size{1};
        int64_t indices_d_size{1};
        bool is_tik{false};
        bool is_gather_v2{false};
        int64_t impl_mode{IMPL_MODE_DEFAULT_VALUE};
        bool is_preprocessed{false};
    };

    struct GatherCompileParams {
        int64_t ub_size;
        int64_t l1_size;
        int64_t core_num;
        int64_t params_d_size;
        int64_t indices_d_size;
        int64_t batch_dims;
        int64_t impl_mode;
        bool is_preprocessed;
    };

    struct GatherShapeInfo {
        std::vector <int64_t> params_shape;
        std::vector <int64_t> indices_shape;
        std::vector <int64_t> indices_ori_shape;
        std::vector <int64_t> y_shape;
    };

    struct GatherV2TilingParams {
        int64_t tiling_mode;
        int64_t params_pre;
        int64_t params_axis;
        int64_t params_row;
        int64_t indices_num;
        int64_t cache_params;
        int64_t need_core_num;
        int64_t tail_process_core;
        int64_t indices_num_each_core;
        int64_t indices_num_remaining;
        int64_t indices_loop_num;
        int64_t indices_row_num_once;
        int64_t indices_row_num_last;
        int64_t row_num_once_ub;
        int64_t row_num_once_tail_ub;
        int64_t inner_loop_num;
        int64_t row_num_last_ub;
        int64_t row_num_last_tail_ub;
        int64_t inner_loop_num_last;
        int64_t params_total;
        int64_t one_row_loop;
        int64_t one_row_tail;
        int64_t params_pre_each_core;
        int64_t params_pre_remaining;
        int64_t indices_row;
        int64_t params_batch_each_core;
        int64_t params_batch_remaining;
        int64_t params_batch;
        int64_t indicesBatch;
        int64_t half_remain_ub_size;
        int64_t half_ub_size;
    };

    bool GatherTIKTiling(const std::string &op_type, const TeOpParas &op_paras, const GatherCompileInfos &op_info,
                         OpRunInfo &run_info);
    void InitGatherCompileParams(GatherCompileParams &params);
    void InitGatherShapeInfo(const std::string &op_type, GatherShapeInfo &params, const TeOpParas &opParas);
    void InitGatherV2Params(GatherV2TilingParams &params);
    void SetGatherV2Params(GatherV2TilingParams &Params, OpRunInfo &run_info);
    void PrintGatherV2Params(const GatherV2TilingParams &params, const std::string &op_type);
    bool CheckAxisAndBatchdims(const std::string &op_type, const GatherShapeInfo &shape_info, int64_t &axis,
        GatherCompileParams &compile_params);
    void GetV2GatherCompileParams(const std::string &op_type, const GatherCompileInfos &compile_info_vec,
        const TeOpParas &op_paras, GatherCompileParams &params);
    bool DoImplModeTiling(const std::string &op_type, GatherV2TilingParams &run_params,
        const GatherCompileParams &compile_params);
    bool DoCacheModeAlignCheck(const std::string &op_type, int64_t axis, GatherV2TilingParams &run_params,
        const GatherCompileParams &compile_params);
    bool DoCacheModeTiling(const std::string &op_type, int64_t axis, GatherV2TilingParams &run_params,
        const GatherCompileParams &compile_params);
    bool DoCacheModeNotAlignCheck(const std::string &op_type, int64_t axis, int64_t one_param_row_size,
        GatherV2TilingParams &params, const GatherCompileParams &compile_params);
    bool DoCacheModeTilingNotAlian(const std::string &op_type, int64_t axis, GatherV2TilingParams &params,
        const GatherCompileParams &compile_params);
    bool DoCpuPreprocessCheck(const std::string &op_type, int64_t axis, int64_t part_ub_size,
        GatherV2TilingParams &params, const GatherCompileParams &compile_params);
    bool DoCacheModeTilingwithCpuPreprocess(const std::string &op_type, int64_t axis, GatherV2TilingParams &params,
        const GatherCompileParams &compile_params);
    bool BlockLessForParamsTiling(GatherV2TilingParams &run_params, int64_t indices_num_per_loop, int64_t res_ub_size,
        int64_t params_d_size, int64_t block_num, const std::string &op_type);
    bool BlockAlignForParamsTiling(GatherV2TilingParams &run_params, int64_t indices_num_per_loop, int64_t res_ub_size,
        int64_t params_d_size, const std::string &op_type);
    bool BlockLessForIndicesTiling(GatherV2TilingParams &run_params, int64_t indices_num_per_loop, int64_t res_ub_size,
        int64_t params_d_size, int64_t block_num, const std::string &op_type);
    bool BlockAlignForIndicesTiling(GatherV2TilingParams &run_params, int64_t indices_num_per_loop, int64_t res_ub_size,
        int64_t params_d_size, const std::string &op_type);
    bool CalcWithBatchDims(GatherV2TilingParams &run_params, int64_t indices_num_per_loop, int64_t res_ub_size,
        int64_t params_d_size);
    bool CalcCacheIndices(GatherV2TilingParams &run_params, int64_t indices_num_per_loop, int64_t res_ub_size,
        int64_t params_d_size, int64_t tiling_mode);
    void CalNeedCoreWithBatchDims(GatherV2TilingParams &run_params, const GatherCompileParams &compile_params);
    void CalNeedCore(GatherV2TilingParams &run_params, const GatherCompileParams &compile_params);
    void ParasPreProcess(GatherV2TilingParams &run_params, const GatherCompileParams &compile_params,
        const GatherShapeInfo &shape_info, int64_t axis);
    bool SmallRowProcess(GatherV2TilingParams &run_params, const GatherCompileParams &compile_params,
        int64_t mode_with_cache, int64_t mode_without_cache);
    bool IndicesCachedProcess(GatherV2TilingParams &run_params, const GatherCompileParams &compile_params,
        int64_t aval_ub_size, int64_t mode_cache_all, int64_t mode_cache_row, int64_t mode_without_cache);
    bool LargeRowProcess(GatherV2TilingParams &run_params, const GatherCompileParams &compile_params,
        int64_t half_ub_params_elem);
    bool TilingWithBatchDims(GatherV2TilingParams &run_params, const GatherCompileParams &compile_params,
        const GatherShapeInfo &shape_info, int64_t axis);
    bool TilingWithoutBatchDims(const std::string &op_type, GatherV2TilingParams &run_params,
        GatherCompileParams &compile_params, const GatherShapeInfo &shape_info, int64_t axis);

}  // namespace optiling
#endif  // AIR_CXX_RUNTIME_V2_OP_IMPL_DYNAMIC_ATOMIC_ADDR_CLEAN_IMPL_H_
